
def restricted_import(name, globals=None, locals=None, fromlist=(), level=0):
	libs = ['math', 'itertools', 'json', 'heapq' , 'sympy', 'numpy',
			'datetime', 'time', 'collections', 'fractions', 'scipy', 'scipy.optimize']
	if name in libs:
		exec(f'import {name}')
		return eval(name)
	raise ImportError(f"Importing {name} is not allowed.")

safe_globals = {
	'print': print,
	'eval': eval,
	'exec': exec,
	'ImportError': ImportError,
	'Exception': Exception,
	'ZeroDivisionError': ZeroDivisionError,
	'ValueError': ValueError,
	'__builtins__': {
		'True': True,
		'False': False,
		'None': None,
		'abs': abs,
		'float': float,
		'int': int,
		'str': str,
		'set': set,
		'list': list,
		'dict': dict,
		'map': map,
		'len': len,
		'max': max,
		'min': min,
		'round': round,
		'sum': sum,
		'pow': pow,
		'range': range,
		'complex': complex,
		'sorted': sorted,
		'reversed': reversed,
		'enumerate': enumerate,
		'zip': zip,
		'any': any,
		'all': all,
		'ord': ord,
		'chr': chr,
		'isinstance': isinstance,
		'filter': filter,
		'next': next,
		'__import__': restricted_import,
	},
}


def execute_code_thread(queue, code, context=None):
	import copy, io, contextlib
	g = copy.deepcopy(safe_globals)
	stderr_buffer = io.StringIO()
	stdout_buffer = io.StringIO()
	try:
		with contextlib.redirect_stderr(stderr_buffer):
			with contextlib.redirect_stdout(stdout_buffer):
				exec(code, g)
		stderr_text = stderr_buffer.getvalue()
		stdout_text = stdout_buffer.getvalue()
		if stdout_text.strip() or stderr_text.strip():
			if not stderr_text.strip():
				stderr_text = None
			queue.put([stdout_text, stderr_text])
			return
	except Exception:
		import traceback
		ex = traceback.format_exc()
		queue.put([None, ex])
		return
	try:
		val = ''
		for l in code.split('\n')[::-1]:
			l = l.strip()
			if l == '' or l[0] == '#':
				continue
			if '=' in l:
				l = l.split('=')[0]
			val = eval(l, g)		
			val = f'{l}: {val}'
			break
	except Exception:
		pass
	queue.put([val, None])
	
def execute_code(code, timeout=60):
	# 解析代码
	import time, ast
	_ = ast.parse(code)
	start_time = time.time()
	from multiprocessing import Process, Queue
	queue = Queue()
	process = Process(target=execute_code_thread, args=(queue, code))
	process.start()
	process.join(timeout=timeout)
	if process.is_alive():
		process.terminate()
		with open('timeout_code.txt', 'a+') as fp:
			fp.write(code+'\n----------\n')
		raise TimeoutError('Refine your approach to implement the code more efficiently')
	assert not queue.empty()
	res, exc = queue.get()
	assert exc is None, exc
	return res

def remove_additional_indent(code):
	lines = code.split('\n')
	line0 = lines[0].lstrip()
	if len(line0) == len(lines[0]):
		return code
	indent = lines[0][:-len(line0)]
	assert len(indent) > 0, lines[0]
	for i, line in enumerate(lines):
		assert line.startswith(indent) or line.strip()=='', [indent, i, line, code]
		lines[i] = line[len(indent):]
	return '\n'.join(lines)

def extract_block(response, exp):
	import re
	if exp == 'python':
		exp = r"```python\n([\s\S]+?)```"
	elif exp == 'json':
		exp = r"```json\n([\s\S]+?)```"
	elif exp == 'markdown':
		exp = r"```markdown\n([\s\S]+?)```"
	else:
		assert False, f'Not supported: {exp}'
	code = re.search(exp, response)
	if code:
		code = code.group(1).rstrip()
		code = remove_additional_indent(code)
		return code
	return None
